{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0bbba98-722c-4a86-a935-354407d684c1",
   "metadata": {},
   "source": [
    "# torch_geometric入门"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "85e9800c-58b3-4bcd-8c6d-f73a1ca97e77",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.x\n",
      "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.tx\n",
      "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.allx\n",
      "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.y\n",
      "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.ty\n",
      "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.ally\n",
      "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.graph\n",
      "Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.test.index\n",
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "from torch_geometric.datasets import Planetoid\n",
    "\n",
    "dataset = Planetoid(root='/tmp/Cora', name='Cora')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8fddc09c-8d0b-44b8-bf35-1e39d3c24363",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv\n",
    "\n",
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = GCNConv(dataset.num_node_features, 16)\n",
    "        self.conv2 = GCNConv(16, dataset.num_classes)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index = data.x, data.edge_index\n",
    "\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "88f90ed8-05bd-4de7-b396-de6372f41138",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1433, 7)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.num_node_features, dataset.num_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "66d39f87-16d6-4a88-a482-2180a0e9f345",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2708, 1433]), torch.Size([2, 10556]))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.x.shape, dataset.edge_index.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7976a715-477c-4a4a-af8c-b3166439c56e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[   0,    0,    0,  ..., 2707, 2707, 2707],\n",
       "        [ 633, 1862, 2582,  ...,  598, 1473, 2706]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1eabb983-7570-439b-8872-2155e8cd9e40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ True,  True,  True,  ..., False, False, False])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.train_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e564b4be-e672-43b1-8a73-22d4883c88db",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = GCN().to(device)\n",
    "data = dataset[0].to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)\n",
    "\n",
    "model.train()\n",
    "for epoch in range(200):\n",
    "    optimizer.zero_grad()\n",
    "    out = model(data)\n",
    "    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7578df63-0d3e-452b-9b76-34be755410cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2708, 7]), torch.Size([140, 7]))"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out.shape, out[dataset.train_mask].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f619e9b8-63f4-414e-a9d7-efdfa86497d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.8050\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "pred = model(data).argmax(dim=1)\n",
    "correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()\n",
    "acc = int(correct) / int(data.test_mask.sum())\n",
    "print(f'Accuracy: {acc:.4f}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00a9fe4a-3d27-4b42-b1bf-8a4b6592660b",
   "metadata": {},
   "source": [
    "# 接下来跟torch_geometric关系不大"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "5e9e03a6-bb14-4334-8620-a03202402614",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse import csr_matrix\n",
    "import pandas as pd\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "from sklearn.metrics import roc_auc_score, roc_curve\n",
    "\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "beb9d946-2d32-4290-9220-2cca436278a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://www.kaggle.com/datasets/alexteboul/heart-disease-health-indicators-dataset\n",
    "path='/Users/qiyusama/Documents/work/datasets/heart_disease_health_indicators_BRFSS2015.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "66a36ccd-f438-4722-8f87-4005aa34d5a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "28c10ea0-75d8-495e-82c2-c2e87df3f949",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['HeartDiseaseorAttack', 'HighBP', 'HighChol', 'CholCheck', 'BMI',\n",
       "       'Smoker', 'Stroke', 'Diabetes', 'PhysActivity', 'Fruits', 'Veggies',\n",
       "       'HvyAlcoholConsump', 'AnyHealthcare', 'NoDocbcCost', 'GenHlth',\n",
       "       'MentHlth', 'PhysHlth', 'DiffWalk', 'Sex', 'Age', 'Education',\n",
       "       'Income'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0c792b21-566a-4cf2-a4a9-7104e328784d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HighBP 2\n",
      "HighChol 2\n",
      "CholCheck 2\n",
      "Smoker 2\n",
      "Stroke 2\n",
      "Diabetes 3\n",
      "PhysActivity 2\n",
      "Fruits 2\n",
      "Veggies 2\n",
      "HvyAlcoholConsump 2\n",
      "AnyHealthcare 2\n",
      "NoDocbcCost 2\n",
      "GenHlth 5\n",
      "DiffWalk 2\n",
      "Sex 2\n",
      "['HighBP', 'HighChol', 'CholCheck', 'Smoker', 'Stroke', 'Diabetes', 'PhysActivity', 'Fruits', 'Veggies', 'HvyAlcoholConsump', 'AnyHealthcare', 'NoDocbcCost', 'GenHlth', 'DiffWalk', 'Sex'] ['BMI', 'MentHlth', 'PhysHlth', 'Age', 'Education', 'Income']\n"
     ]
    }
   ],
   "source": [
    "cate_cols = []\n",
    "num_cols = []\n",
    "label_col = 'HeartDiseaseorAttack'\n",
    "for c in df.columns[1:]:\n",
    "    num_unique = df[c].nunique()\n",
    "    if num_unique<=5:\n",
    "        cate_cols.append(c)\n",
    "        print(c, num_unique)\n",
    "    else:\n",
    "        num_cols.append(c)\n",
    "print(cate_cols, num_cols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "cdce3dc5-5552-4f6b-ae5a-1471da11515a",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.sample(n=10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ba1e857d-7bc8-4b34-adf0-d9e10bdf9b15",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeartDiseaseorAttack\n",
       "0.0    9047\n",
       "1.0     953\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[label_col].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "86b5780f-2955-4386-bb83-05870674e72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建 OneHotEncoder 对象\n",
    "onehot_encoder = OneHotEncoder(sparse_output=False, drop='first')  # drop='first' 可以避免多重共线性\n",
    "\n",
    "# 对类别型变量进行 One-Hot 编码\n",
    "encoded_data = onehot_encoder.fit_transform(df[cate_cols])\n",
    "\n",
    "# 获取新生成的列名\n",
    "encoded_columns = onehot_encoder.get_feature_names_out(cate_cols)\n",
    "\n",
    "# 将编码后的数据转换为 DataFrame\n",
    "adj_bi = pd.DataFrame(encoded_data, columns=encoded_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c3414dee-1077-432e-8ec8-68fad74cddd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_adj = csr_matrix(adj_bi)\n",
    "adj = sp_adj*sp_adj.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "b2af8d87-7ec7-4f0a-993d-e009c40e77aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = adj.A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "fdd03dfc-62aa-4dd9-a4a5-443533831dc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def graph_normalize(adj):\n",
    "    degree=torch.sum(adj,dim=1)\n",
    "    d_inv_sqrt = torch.diag(torch.pow(degree,-0.5))\n",
    "    norm_adj = torch.mm(torch.mm(d_inv_sqrt, adj), d_inv_sqrt)\n",
    "    return norm_adj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "05addbbb-1b87-46ec-a8e8-da0152c4ddaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = torch.from_numpy(A).type(torch.FloatTensor)\n",
    "A = graph_normalize(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "ef005a09-8774-46ce-b1bb-ca98395f15fb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10000, 6)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "a245ea27-f859-4962-893c-5b0af86db18e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "import numpy as np\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from sklearn.neighbors import kneighbors_graph\n",
    "\n",
    "X = df[num_cols].values\n",
    "sc_m = MinMaxScaler()\n",
    "X = sc_m.fit_transform(X)\n",
    "\n",
    "X = torch.from_numpy(X).float()\n",
    "y = torch.from_numpy(df[label_col].values).long()\n",
    "\n",
    "indices = np.arange(X.shape[0])\n",
    "train_idx, tmp_idx, y_train, y_tmp = train_test_split(indices, y, test_size=0.3, random_state=42)\n",
    "val_idx, test_idx, y_val, y_test = train_test_split(indices, y, test_size=0.4, random_state=42)\n",
    "\n",
    "X_train, X_val, X_test = X[train_idx], X[val_idx], X[test_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "889af070-2fe9-4126-9f23-8bab6af5519d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.parameter import Parameter\n",
    "import math\n",
    "\n",
    "class GraphConvolution(nn.Module):\n",
    "    def __init__(self, in_features, out_features, bias=True):\n",
    "        super(GraphConvolution, self).__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.weight = Parameter(torch.FloatTensor(in_features, out_features))\n",
    "        if bias:\n",
    "            self.bias = Parameter(torch.FloatTensor(out_features))\n",
    "        else:\n",
    "            self.register_parameter('bias', None)\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        stdv = 1. / math.sqrt(self.weight.size(1))\n",
    "        self.weight.data.uniform_(-stdv, stdv)\n",
    "        if self.bias is not None:\n",
    "            self.bias.data.uniform_(-stdv, stdv)\n",
    "\n",
    "    def forward(self, inputs, adj):\n",
    "        support = torch.spmm(inputs, self.weight)\n",
    "        output = torch.spmm(adj, support)\n",
    "        if self.bias is not None:\n",
    "            return F.elu(output + self.bias)\n",
    "        else:\n",
    "            return F.elu(output)\n",
    "\n",
    "class GCN(nn.Module):\n",
    "    def __init__(self, in_size, hidden_size, num_classes, dropout):\n",
    "        super(GCN, self).__init__()\n",
    "        # self.fc_trans = nn.Linear(in_size, hidden_size)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        self.gcn_layers = GraphConvolution(in_size, hidden_size)\n",
    "        self.gcn_layers2 = GraphConvolution(hidden_size, hidden_size)\n",
    "        self.fc_class = nn.Linear(hidden_size, num_classes)\n",
    "\n",
    "    def forward(self, x, adj):\n",
    "        # h = self.fc_trans(x)\n",
    "        # h = self.dropout(h)\n",
    "        h = self.gcn_layers(x, adj)\n",
    "        h = self.dropout(h)\n",
    "        h = self.gcn_layers2(h, adj)\n",
    "        h = self.dropout(h)\n",
    "        h = self.fc_class(h)\n",
    "        return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "c5f15d37-2811-43f1-a4bd-70025a16d8b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_size = X.shape[1]\n",
    "hidden_size = 64\n",
    "num_classes = len(np.unique(y.cpu().numpy()))\n",
    "epochs = 500\n",
    "early_stop = 30\n",
    "model = GCN(input_size, hidden_size, num_classes, 0.5)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=5e-3, weight_decay=1e-3)\n",
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "best_val_loss=10.0\n",
    "patience=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "5d7e166f-3da2-4289-8048-f2826da749f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10000, 6])"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "c0b1a9d1-fd15-4a35-831a-bb0f768654a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0471, -0.9360],\n",
       "        [ 1.3413, -1.1899],\n",
       "        [ 1.1583, -1.0367],\n",
       "        ...,\n",
       "        [ 1.1391, -0.9543],\n",
       "        [ 1.1287, -1.0151],\n",
       "        [ 1.0615, -0.8035]], grad_fn=<IndexBackward0>)"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "3a3935e0-e220-4b8e-af06-613280523962",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 20, Train Acc: 0.9011, Train Loss: 0.3162, Val Acc: 0.9033, Val Loss: 0.3111, Test Acc: 0.9067, Test Loss: 0.3036\n",
      "[auc] train:0.6095333570362984,val:0.6293712622471053,test:0.617583272906286\n",
      "Epoch: 40, Train Acc: 0.9011, Train Loss: 0.3187, Val Acc: 0.9033, Val Loss: 0.3109, Test Acc: 0.9067, Test Loss: 0.3034\n",
      "[auc] train:0.5930334129337551,val:0.6294293167069602,test:0.6176084046446408\n",
      "Epoch: 60, Train Acc: 0.9011, Train Loss: 0.3177, Val Acc: 0.9033, Val Loss: 0.3109, Test Acc: 0.9067, Test Loss: 0.3035\n",
      "[auc] train:0.6017179075291125,val:0.6294522203842727,test:0.6176764081719543\n",
      "Epoch: 80, Train Acc: 0.9011, Train Loss: 0.3177, Val Acc: 0.9033, Val Loss: 0.3111, Test Acc: 0.9067, Test Loss: 0.3038\n",
      "[auc] train:0.6017547906869339,val:0.6294592187301182,test:0.6176963657288832\n",
      "Epoch: 100, Train Acc: 0.9011, Train Loss: 0.3174, Val Acc: 0.9033, Val Loss: 0.3108, Test Acc: 0.9067, Test Loss: 0.3033\n",
      "[auc] train:0.6085230792351028,val:0.629466217075964,test:0.6176971048976584\n",
      "Epoch: 120, Train Acc: 0.9011, Train Loss: 0.3171, Val Acc: 0.9033, Val Loss: 0.3108, Test Acc: 0.9067, Test Loss: 0.3032\n",
      "[auc] train:0.6016339468002829,val:0.6294713067820333,test:0.6177056053385725\n",
      "Epoch: 140, Train Acc: 0.9011, Train Loss: 0.3182, Val Acc: 0.9033, Val Loss: 0.3109, Test Acc: 0.9067, Test Loss: 0.3035\n",
      "[auc] train:0.5955920960996404,val:0.6294856215803537,test:0.6178013276949539\n",
      "Epoch: 160, Train Acc: 0.9011, Train Loss: 0.3175, Val Acc: 0.9033, Val Loss: 0.3112, Test Acc: 0.9067, Test Loss: 0.3039\n",
      "[auc] train:0.6051089817132845,val:0.6295234762692454,test:0.6178648962096165\n",
      "Epoch: 180, Train Acc: 0.9011, Train Loss: 0.3159, Val Acc: 0.9033, Val Loss: 0.3108, Test Acc: 0.9067, Test Loss: 0.3032\n",
      "[auc] train:0.6188870174949876,val:0.6295272935487977,test:0.6178249810957586\n",
      "Early stopping at epoch 192\n"
     ]
    }
   ],
   "source": [
    "def auc(output, y):\n",
    "    softmax_output = F.softmax(output, dim=1)\n",
    "    positive_probs = softmax_output[:, 1]\n",
    "    y = y.cpu().numpy()\n",
    "    positive_probs = positive_probs.cpu().detach().numpy()\n",
    "    auc = roc_auc_score(y, positive_probs)\n",
    "    return auc\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    output = model(X,A)\n",
    "    train_output = output[train_idx]\n",
    "    loss = criterion(train_output, y_train)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    train_acc = (train_output.max(1)[1] == y_train).sum().item() / len(y_train)\n",
    "    train_auc = auc(train_output, y_train)\n",
    "    # 在验证集上评估模型\n",
    "    model.eval()\n",
    "    output = model(X, A)\n",
    "    val_output = output[val_idx]\n",
    "    val_loss = criterion(val_output, y_val).item()\n",
    "    val_acc = (val_output.max(1)[1] == y_val).sum().item() / len(y_val)\n",
    "    val_auc = auc(val_output, y_val)\n",
    "\n",
    "    test_output = output[test_idx]\n",
    "    test_loss = criterion(test_output, y_test).item()\n",
    "    test_acc = (test_output.max(1)[1] == y_test).sum().item() / len(y_test)\n",
    "    test_auc = auc(test_output, y_test)\n",
    "\n",
    "    num = epoch + 1\n",
    "    if num % 20 == 0:\n",
    "        print(f'Epoch: {epoch+1}, Train Acc: {train_acc:.4f}, Train Loss: {loss:.4f}, Val Acc: {val_acc:.4f}, Val Loss: {val_loss:.4f}, Test Acc: {test_acc:.4f}, Test Loss: {test_loss:.4f}')\n",
    "        print(f'[auc] train:{train_auc},val:{val_auc},test:{test_auc}')\n",
    "    # 早停\n",
    "    if val_loss < best_val_loss:\n",
    "        best_val_loss = val_loss\n",
    "        patience = 0\n",
    "        torch.save(model.state_dict(), 'model.pth')\n",
    "    else:\n",
    "        patience += 1\n",
    "        if patience >= early_stop:\n",
    "            print(f'Early stopping at epoch {epoch+1}')\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a20ffed2-681f-444e-b246-53df39802068",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 在测试集上评估模型\n",
    "model.load_state_dict(torch.load('model.pth'))\n",
    "model.eval()\n",
    "output = model(A, X)\n",
    "test_output = output[test_indices]\n",
    "test_loss = criterion(test_output, y_test).item()\n",
    "test_acc = (test_output.max(1)[1] == y_test).sum().item() / len(y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b1a758e-b42e-4970-8a86-d262bca1bcbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 将编码后的 DataFrame 与原始 DataFrame 中的非类别列合并\n",
    "non_cate_cols = df.drop(columns=cate_cols)\n",
    "new_df = pd.concat([non_cate_cols, encoded_df], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "50ab3d69-11d0-44ca-b134-7a38ae012006",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(253680, 19)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded_df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e06a6df8-3db9-48f2-9601-1213c3d9a197",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeartDiseaseorAttack\n",
       "0.0    229787\n",
       "1.0     23893\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[label_col].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "558ab044-16ea-44d4-83fd-a35fafe19341",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>BMI</th>\n",
       "      <th>MentHlth</th>\n",
       "      <th>PhysHlth</th>\n",
       "      <th>Age</th>\n",
       "      <th>Education</th>\n",
       "      <th>Income</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>40.0</td>\n",
       "      <td>18.0</td>\n",
       "      <td>15.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>25.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>7.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>28.0</td>\n",
       "      <td>30.0</td>\n",
       "      <td>30.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>8.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>27.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>6.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>24.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>4.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>253675</th>\n",
       "      <td>45.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>7.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>253676</th>\n",
       "      <td>18.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>4.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>253677</th>\n",
       "      <td>28.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>253678</th>\n",
       "      <td>23.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>7.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>253679</th>\n",
       "      <td>25.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>253680 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         BMI  MentHlth  PhysHlth   Age  Education  Income\n",
       "0       40.0      18.0      15.0   9.0        4.0     3.0\n",
       "1       25.0       0.0       0.0   7.0        6.0     1.0\n",
       "2       28.0      30.0      30.0   9.0        4.0     8.0\n",
       "3       27.0       0.0       0.0  11.0        3.0     6.0\n",
       "4       24.0       3.0       0.0  11.0        5.0     4.0\n",
       "...      ...       ...       ...   ...        ...     ...\n",
       "253675  45.0       0.0       5.0   5.0        6.0     7.0\n",
       "253676  18.0       0.0       0.0  11.0        2.0     4.0\n",
       "253677  28.0       0.0       0.0   2.0        5.0     2.0\n",
       "253678  23.0       0.0       0.0   7.0        5.0     1.0\n",
       "253679  25.0       0.0       0.0   9.0        6.0     2.0\n",
       "\n",
       "[253680 rows x 6 columns]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[num_cols]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "143e23e5-139a-425f-bc3c-eed2da3b5460",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0c0c654-e94b-426d-bcc0-974ea4d03905",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11",
   "language": "python",
   "name": "myenv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
